package com.wenhai.leetcode.sort;

import java.util.PriorityQueue;
import java.util.Random;

/**
 * https://leetcode-cn.com/problems/kth-largest-element-in-an-array/
 *
 * @author wenhai
 * @date   2021/3/1
 */
public class _215数组中的第K个最大元素 {
    private static final Random RANDOM = new Random();

    public int findKthLargestR(int[] nums, int k) {

        return findKthLargestR(nums, 0, nums.length - 1, nums.length  - k);

    }

    private int findKthLargestR(int[] nums, int l, int r, int k) {
        int p = partition(nums, l, r);
        if (k == p) {
            return nums[p];
        } else if (p < k) {
            return findKthLargestR(nums, p + 1, r, k);
        } else {
            return findKthLargestR(nums, l, p - 1, k);
        }

    }

    private int partition(int[] nums, int l, int r) {
        int p = l + RANDOM.nextInt(r - l + 1);
        swap(nums, l, p);

        int j = r;
        int i = l + 1;

        while (true){
            while (i <= j && nums[i] < nums[l]) {
                i++;
            }
            while (j >= i && nums[j] > nums[l]) {
                j--;
            }
            if (i >= j) {
                break;
            }
            swap(nums, i, j);
            i++;
            j--;
        }
        swap(nums, l, j);
        return j;
    }

    private void swap(int[] nums, int srcPos, int destPos) {
        int temp = nums[srcPos];
        nums[srcPos] = nums[destPos];
        nums[destPos] = temp;
    }


    private int findKthLargest(int[] nums, int k) {
        k = nums.length - k;

        int l = 0;
        int r = nums.length - 1;

        while (l <= r) {
            int p = partition(nums, l, r);
            if (k == p) {
                return nums[p];
            } else if (p < k) {
                l = p + 1;
            } else {
                r = p - 1;
            }
        }
       return -1;

    }

    public int findKthLargest2(int[] nums, int k) {

        PriorityQueue<Integer> pq = new PriorityQueue<>(k);
        for (int i = 0; i < k; i++) {
            pq.add(nums[i]);
        }
        for (int i = k; i < nums.length; i++) {
            if (!pq.isEmpty() && pq.peek() < nums[i]) {
                pq.remove();
                pq.add(nums[i]);
            }
        }

        return pq.peek();
    }
    public static void main(String[] args) {
        int[] arr = {3, 2, 1, 5, 6, 4};
        System.out.println(new _215数组中的第K个最大元素().findKthLargest(arr, 2));


    }
}
